# encoding: utf-8
"""
中心控制组件
"""
import os
import copy
import gc
import logging
import traceback
import uuid
import time
import threading
from concurrent.futures import ThreadPoolExecutor

import psycopg2

from .source import SourceDB, LogCache, write_to_json
from ..model.data import create_preheating_data
from ..config.base import PRINT_TRACEBACK_OPTION, ERROR_INTERSECT_CORRECTION, ERROR_CORRECTION_MODEL_PATH, ASSIST_ERROR_CORRECTION_MODEL_PATH, ERROR_IGNORE_DEL
from ..config.db import POSTGRESQL_INFO, get_init_params
from ..config.error import DataCompleteError, ControlAssemblyError
from ..config.interface import PREDICT_REQUEST_WAIT_COUNT, PREDICT_PER_WAIT_TIME, UPDATE_INTERVAL_SECONDS, \
    HEARTBEAT_DELETE_TIME
from ..model.base import create_model_with_name
from ..model.rex.predict import predict_contract_elements
from ..model.error_correction import (create_error_correction_model,
                                             new_correct_text_error, intersect_correct_text_error, del_ignore_errors)
from ..model.rex.train import DEFAULT_TRAIN_CONFIG, run_train, \
    save_model, switch_params, get_save_model_path, get_train_data_path, rm_train_data
from ..embedding.embedding import Embedding

POOL_COUNT = 1
split_ratio = 0.2


class Control(object):
    def __init__(self):
        """
        初始化参数
        """
        self.server_id = str(uuid.uuid1())
        self.pid = os.getpid()
        '''
        从db获取基本参数
        '''
        init_params = get_init_params()
        self.default_model_id = init_params['default_model_id']
        self.now_model_id = None
        '''
        构建基本参数
        '''
        self.train_task_queue = []  # 模型训练队列
        self.train_task_poll = None  # 训练任务池
        self.is_predict = False  # 是否正在预测中
        self.model_status = 'FREE'  # 模型应该的状态
        self.is_rag = False  # 模型的rag运行
        self.is_correct_error = False  # 错字审查

        '''
        构建进程信息
        '''
        self.create_server_info()

        '''
        加载合同要素提取模型
        '''
        self.contract_model = None
        self.load_model(self.default_model_id)  # 合同要素模型

        '''
        加载错字检测模型
        '''
        self.correct_error_model = create_error_correction_model(ERROR_CORRECTION_MODEL_PATH)
        # 进行模型交叉验证
        if ERROR_INTERSECT_CORRECTION:
            self.assist_correct_error_model = create_error_correction_model(ASSIST_ERROR_CORRECTION_MODEL_PATH)
        # self.correct_error_model = None

        '''
        线程池
        用于异步训练
        '''
        self.parallel_pool = ThreadPoolExecutor(POOL_COUNT)

        '''
        初始化
        '''
        self._init_run()

        '''
        启动数据库刷新监听
        '''
        self.start_update_info_listening()
        self.embedding = Embedding()


    def create_server_info(self):
        SourceDB.create_server_info(self.server_id, self.default_model_id, self.now_model_id, self.pid)

    def update_server_info_with_load(self, load_model_label, load_model_error):
        SourceDB.update_server_info_with_load(self.server_id, self.default_model_id, self.now_model_id,
                                              load_model_label, load_model_error)

    def correct_error_text(self, text_infos):
        # 校验文本差异
        # result = correct_text_error(text_infos, self.correct_error_model)
        if self.is_correct_error:
            raise ControlAssemblyError('错字审查频繁, 请稍后再试')
        try:
            self.is_correct_error = True
            if not ERROR_INTERSECT_CORRECTION:
                result = new_correct_text_error(text_infos, self.correct_error_model, 'other')
            else:
                # 交叉辅助验证模块
                result = intersect_correct_text_error(text_infos, self.correct_error_model, self.assist_correct_error_model, True, 'other')
            # 删除忽略的错误
            if ERROR_IGNORE_DEL:
                result = del_ignore_errors(text_infos, result)
        finally:
            self.is_correct_error = False
        return result

    def start_update_info_listening(self):
        timer = threading.Timer(UPDATE_INTERVAL_SECONDS, self.start_update_info_listening)
        timer.start()

        # 在定时器触发时运行
        self._update_model_info()


    def _update_model_info(self):
        # 进行心跳
        SourceDB.update_server_heartbeat(self.server_id)
        # 删除超过心跳
        SourceDB.delete_server_info(HEARTBEAT_DELETE_TIME)

        model_info = SourceDB.get_default_info()
        model_status = model_info['model_status']
        default_model_id = model_info['default_model_id']
        # 模型状态切分
        if self.model_status != model_status:
            # 判断是否加载模型
            if model_status == 'FREE' and self.contract_model is None:
                self.default_model_id = default_model_id
                self.load_model(self.default_model_id)
            # 判断是否释放模型
            if model_status == 'TRAINING' and self.contract_model is not None:
                self._release_model()
            self.model_status = model_status

        # 切换模型
        if self.default_model_id != default_model_id:
            self.default_model_id = default_model_id
            self._switch_model(self.default_model_id)

        # import datetime
        # print(datetime.datetime.now(), model_info)

    def _init_run(self):
        """
        初始化一些运行组件
        :return:
        """
        '''
        删除一些生成的训练和测试文件
        '''
        rm_train_data(None)
        logging.info('删除遗留的训练数据文件')

    def judge_model_work(self):
        """
        检验模型是否在运行中, 因为显存不足 只能单次访问
        :return:
        """
        option = True if self.model_status != 'FREE' or self.is_predict else False
        return option

    def train_model_task(self, train_task_id):
        """
        启动模型训练任务束
        :param train_task_id:
        :return:
        """
        option = {
            'success': False,
            'error': None,
            'result_msg': None
        }
        # 添加训练任务
        if train_task_id in self.train_task_queue:
            option['success'] = False
            result_msg = '有相同任务在训练中或者训练队列中'
            logging_msg = '添加训练任务: %s失败' % train_task_id + ' ' + result_msg
            logging.error(logging_msg)
            option['result_msg'] = result_msg
        else:
            option['success'] = True
            self.train_task_queue.append(train_task_id)
            logging_msg = '添加训练任务：%s' % train_task_id
            logging.info(logging_msg)
            option['result_msg'] = logging_msg
            # 开始任务
            # 以单核并行方式训练模型
            SourceDB.update_train_queue_task_start(train_task_id, None)
            self.parallel_pool.submit(self._run_train_schedule, train_task_id)

        return option

    # def get_contract_element(self, infos):
    #     """
    #     :param infos: 要素集合
    #     :return:
    #     """
    #     if self.contract_model is None:
    #         raise Exception('没有加载模型')

    def _release_model(self):
        """
        释放内存
        :return:
        """
        if self.contract_model is not None:
            del self.contract_model
            gc.collect()
        self.contract_model = None

    def wait_resource_release(self, log_cache=None):
        limit_wait_count = 10
        wait_time = 30
        wait_count = 0
        resource_free = False
        while wait_count <= limit_wait_count:
            # 判断是否有模型在运行或者训练
            if self.judge_model_work():
                time.sleep(wait_time)
                wait_count += 1
                if log_cache is not None:
                    log_text = '有模型在运行或者训练, 需要等待%s秒, 等待%s/%s' % (wait_time, wait_count, limit_wait_count)
                    log_cache.add_log_line(log_text, True)
            else:
                resource_free = True
                break
        return resource_free

    def _run_train_schedule(self, train_task_id):
        """
        :param train_task_id:
        :return:
        """
        '''
        创建数据库连接
        '''
        cnx = psycopg2.connect(**POSTGRESQL_INFO)

        log_cache = LogCache(cnx, train_task_id)

        # 开始运行训练程序
        SourceDB.update_train_task_start(train_task_id, cnx)

        '''
        验证资源是否空值
        '''
        resource_free = self.wait_resource_release()

        if not resource_free:
            error_msg = '等待资源空置超时'
            log_cache.add_log_line(error_msg, True)
            logging.error(error_msg)
            log_cache.update_log_to_db()
            # 更新训练完成结果
            training_infos = None
            trained_infos = {
                'return_code': -1
            }
            SourceDB.update_train_task_end(train_task_id, training_infos, trained_infos, cnx)
            self.train_task_queue.remove(train_task_id)
            raise ControlAssemblyError(error_msg)

        try:
            self.model_status = 'TRAINING'
            SourceDB.update_model_status(self.model_status)
            training_infos, trained_infos = self._run_train_model(train_task_id, cnx, log_cache)
        except:
            training_infos = None
            trained_infos = {
                'return_code': -1,
                'model_id': None,
                'metric_span': None
            }
            if PRINT_TRACEBACK_OPTION:
                traceback.print_exc()
            logging.error('模型训练出错')
            error_msg = traceback.format_exc()
            logging.error(error_msg)
            log_cache.add_log_line(error_msg)
        finally:
            self.train_task_queue.remove(train_task_id)
            self.model_status = 'FREE'
            SourceDB.update_model_status(self.model_status)

        # 更新完整日志信息
        log_cache.update_log_to_db()

        # 更新训练完成结果
        SourceDB.update_train_task_end(train_task_id, training_infos, trained_infos, cnx)
        '''
        加载老模型
        '''
        self.load_model(self.default_model_id)

    def _run_train_model(self, train_task_id, cnx, log_cache: LogCache):
        """
        1、 提取训练参数信息
        2、 提取训练数据
        3、 清洗数据集
        4、 开始训练
        :param train_task_id: 训练任务id
        :param train_task_id: cnx
        :param log_cache: 存储的logging信息
        :return:
        """

        def _add_log(logging_text):
            logging.info(logging_text)
            log_cache.add_log_line(logging_text, True)

        '''
        生成模型ID
        '''
        save_model_name = str(str(uuid.uuid1()))

        '''
        释放现有模型
        '''
        self._release_model()
        _add_log('释放现有模型内存空间')

        '''
        获取模型训练参数
        '''
        train_info_result = SourceDB.get_train_info(train_task_id)
        if train_info_result['error_msg'] is not None:
            raise DataCompleteError(train_info_result['error_msg'])

        base_on_model_id = train_info_result['base_on_model_id']
        training_params = train_info_result['training_params']
        available_params = switch_params(training_params)
        if base_on_model_id != '0':
            available_params['bert_model_dir'] = get_save_model_path(base_on_model_id)
        _add_log('获取模型参数')

        '''
        获取数据集
        '''
        _add_log('开始获取模型训练数据')
        all_mark_data = SourceDB.get_dataset_with_task(train_task_id, cnx)
        _add_log('模型训练数据获取完成')
        split_index = int((1 - split_ratio) * len(all_mark_data))
        _add_log('拆分数据, 训练数据: %s, 测试数据: %s' % (split_index, len(all_mark_data) - split_index))

        all_train_data, all_test_data = all_mark_data, all_mark_data[split_index:]
        if len(all_train_data) == 0:
            all_train_data = all_mark_data
        if len(all_test_data) == 0:
            all_test_data = all_mark_data

        train_data_path = get_train_data_path(save_model_name)
        write_to_json(all_train_data, train_data_path + '//' + 'train.json')
        write_to_json(all_test_data, train_data_path + '//' + 'dev.json')
        available_params['data_path'] = train_data_path
        available_params['run_name'] = save_model_name
        actual_model_params = copy.deepcopy(DEFAULT_TRAIN_CONFIG)
        actual_model_params.update(available_params)
        '''
        训练模型
        '''
        _add_log('开始训练模型')
        return_code, metric_span = run_train(actual_model_params, log_cache)
        # if return_code != 0:
        #     raise Exception('模型训练出错')
        '''
        保存模型
        '''
        if return_code == 0:
            logging.info('保存现有模型')
            save_model(save_model_name, save_model_name)
        trained_infos = {
            'return_code': return_code,
            'model_id': save_model_name,
            'metric_span': metric_span
        }
        training_infos = {
            'base_on_model_id': base_on_model_id,
            'training_params': training_params,
        }
        return training_infos, trained_infos

    def switch_model(self, model_id):
        """
        1、将数据库基本模型转换
        2、切换临时数据库默认数据库
        3、等待进程切换结果
        :param model_id:
        :return:
        """
        SourceDB.update_model_default(model_id)
        time.sleep(10)
        option = self.check_switch_model(model_id)
        return option

    def check_switch_model(self, model_id):
        option = {
            'success': False,
            'error_msg': None
        }
        re_connection = []

        def _check_switch():
            # 等待服务器获取线上服务器 并更新模型
            server_infos = SourceDB.get_model_server()
            # 检验所有server是否启动完成
            accuracy_infos = server_infos[
                (server_infos['default_model_id'] == model_id) & (server_infos['now_model_id'] == model_id) & (server_infos[
                    'load_model_label'] == 1)]
            if len(accuracy_infos) == len(server_infos):
                option['success'] = True
            else:
                server_error_infos = server_infos[server_infos['load_model_label'] == 2]
                if len(server_error_infos) > 0:
                    option['success'] = False
                    option['error_msg'] = server_error_infos['load_model_error'].to_list()[0]

                else:
                    if len(re_connection) < 1:
                        re_connection.append(1)
                        time.sleep(30)
                        _check_switch()
                    else:
                        option['success'] = False
                        option['error_msg'] = '重启超时'
        _check_switch()
        return option

    def _switch_model(self, model_name):
        """
        :param model_name:
        :return:
        """
        # 先释放模型
        self._release_model()
        # 再重新加载
        self.load_model(model_name)

    def load_model(self, model_name):
        """
        加载模型
        :return:
        """
        self.now_model_id = model_name
        try:
            self.contract_model = create_model_with_name(model_name)
            load_model_label = 1
            load_model_error = None
            # 预热模型
            self.preheating_model()
        except:
            load_model_label = 2
            load_model_error = traceback.format_exc()
            traceback.print_exc()
            # raise ControlAssemblyError('模型加载出错')
        '''
        更新数据库server表信息
        '''
        self.update_server_info_with_load(load_model_label, load_model_error)

    def judge_resource_free(self):
        if self.model_status == 'TRAINING':
            free_option = False
            free_msg_error = '有模型正在训练, 无法进行要素提取'
            return free_option, free_msg_error

        '''
        等待预测资源释放
        '''
        free_option = True
        free_msg_error = None
        be_wait_count = 0

        while self.is_predict:
            time.sleep(PREDICT_PER_WAIT_TIME)
            be_wait_count += 1
            if be_wait_count >= PREDICT_REQUEST_WAIT_COUNT:
                free_option = False
                free_msg_error = '服务器繁忙, 请稍后再试'
        return free_option, free_msg_error

    def predict_elements(self, model_name, infos, elements):
        free_option, free_msg_error = self.judge_resource_free()

        if not free_option:
            raise ControlAssemblyError(free_msg_error)

        if model_name != self.default_model_id and model_name is not None:
            self.load_model(model_name)

        '''
        校验是否有模型
        '''
        if self.contract_model is None:
            self.load_model(self.default_model_id)

        if isinstance(elements, list):
            default_schema = dict((_, None) for _ in elements)
        else:
            default_schema = elements

        self.is_predict = True

        try:
            time_1 = time.time()
            predict_elements = predict_contract_elements(infos, self.contract_model, default_schema,
                                                         is_relative_index=True, is_switch_data=True)
            time_2 = time.time()
            # print('time_2 - time_1', time_2 - time_1)
            # print(infos)
            # print(default_schema)
            result = {
                'success': True,
                'predict_elements': predict_elements,
                'error_msg': None
            }
        except:
            if PRINT_TRACEBACK_OPTION:
                traceback.print_exc()
            error_msg = traceback.format_exc()
            result = {
                'success': False,
                'predict_elements': None,
                'error_msg': error_msg
            }
        finally:
            self.is_predict = False

        # 将模型切换回原始模型
        if model_name != self.default_model_id and model_name is not None:
            self.load_model(self.default_model_id)
        return result

    def preheating_model(self):
        """
        预热模型
        :return:
        """
        infos, default_schema = create_preheating_data()
        result = self.predict_elements(None, infos, list(default_schema.keys()))
        return result

    def query_similar_text(self, queries, sorted_texts, task, top_n):
        return self.embedding.query_similar_text(queries, sorted_texts, task, top_n)


def get_available_params(available_name, provide_name, available_params, provide_params):
    if provide_name in provide_params:
        pass
    else:
        available_params[available_name] = provide_params[provide_name]
